package org.apache.flink.connector.scylladb.serialization;

import com.datastax.driver.core.BoundStatement;
import com.datastax.driver.core.PreparedStatement;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.data.StringData;
import org.apache.flink.table.data.ArrayData;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.logical.ArrayType;
import static org.apache.flink.table.types.logical.LogicalTypeRoot.*;

import java.math.BigDecimal;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.Date;
import java.util.Set;
import java.util.HashSet;
import java.util.List;

/**
 * ScyllaDB 行数据序列化器，将 Flink RowData 转换为 Cassandra BoundStatement
 */
public class ScyllaDBRowDataSerializer {

    private final RowType rowType;
    private final List<String> primaryKeys;

    public ScyllaDBRowDataSerializer(RowType rowType, List<String> primaryKeys) {
        this.rowType = rowType;
        this.primaryKeys = primaryKeys != null ? primaryKeys : new java.util.ArrayList<>();
    }
    
    // 保持向后兼容的构造函数
    public ScyllaDBRowDataSerializer(RowType rowType) {
        this(rowType, null);
    }

    public BoundStatement serializeInsert(RowData row, PreparedStatement preparedStatement) {
        BoundStatement boundStatement = preparedStatement.bind();
        
        for (int i = 0; i < rowType.getFieldCount(); i++) {
            LogicalType fieldType = rowType.getTypeAt(i);
            Object value = convertValue(row, i, fieldType);
            // 假设 setObject 用于多类型，需根据类型分支调用 setInt/setString/setLong 等。这里用 setInt 作为示例，实际应根据字段类型判断。
            // 例如：boundStatement.setInt(i, (Integer) value);
            // TimestampData.toInstant(ZoneOffset) 改为 TimestampData.toInstant()
            switch (fieldType.getTypeRoot()) {
                case CHAR:
                case VARCHAR:
                    boundStatement.setString(i, (String) value);
                    break;
                case BOOLEAN:
                    boundStatement.setBool(i, (Boolean) value);
                    break;
                case TINYINT:
                    boundStatement.setInt(i, (Integer) value);
                    break;
                case SMALLINT:
                    boundStatement.setInt(i, (Integer) value);
                    break;
                case INTEGER:
                    boundStatement.setInt(i, (Integer) value);
                    break;
                case BIGINT:
                    boundStatement.setLong(i, (Long) value);
                    break;
                case FLOAT:
                    boundStatement.setFloat(i, (Float) value);
                    break;
                case DOUBLE:
                    boundStatement.setDouble(i, (Double) value);
                    break;
                case DECIMAL:
                    boundStatement.setDecimal(i, (BigDecimal) value);
                    break;
                case TIMESTAMP_WITHOUT_TIME_ZONE:
                case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
                    boundStatement.setTimestamp(i, (Date) value);
                    break;
                case BINARY:
                case VARBINARY:
                    boundStatement.setBytes(i, (java.nio.ByteBuffer) value);
                    break;
                case ARRAY:
                    // 处理 Set 类型
                    @SuppressWarnings("unchecked")
                    Set<Object> setData = (Set<Object>) value;
                    boundStatement.setSet(i, setData);
                    break;
                default:
                    boundStatement.setString(i, (String) value);
                    break;
            }
        }
        
        return boundStatement;
    }

    public BoundStatement serializeDelete(RowData row, PreparedStatement preparedStatement) {
        BoundStatement boundStatement = preparedStatement.bind();
        
        // 对于删除操作，通常只需要主键字段
        // 这里简化处理，假设第一个字段是主键
        LogicalType fieldType = rowType.getTypeAt(0);
        Object value = convertValue(row, 0, fieldType);
        // 假设 setObject 用于多类型，需根据类型分支调用 setInt/setString/setLong 等。这里用 setInt 作为示例，实际应根据字段类型判断。
        // 例如：boundStatement.setInt(0, (Integer) value);
        // TimestampData.toInstant(ZoneOffset) 改为 TimestampData.toInstant()
        switch (fieldType.getTypeRoot()) {
            case CHAR:
            case VARCHAR:
                @SuppressWarnings("unchecked")
                Set<Object> setData = (Set<Object>) value;
                boundStatement.setSet(0, setData);
                break;
            case BOOLEAN:
                boundStatement.setBool(0, (Boolean) value);
                break;
            case TINYINT:
                boundStatement.setInt(0, (Integer) value);
                break;
            case SMALLINT:
                boundStatement.setInt(0, (Integer) value);
                break;
            case INTEGER:
                boundStatement.setInt(0, (Integer) value);
                break;
            case BIGINT:
                boundStatement.setLong(0, (Long) value);
                break;
            case FLOAT:
                boundStatement.setFloat(0, (Float) value);
                break;
            case DOUBLE:
                boundStatement.setDouble(0, (Double) value);
                break;
            case DECIMAL:
                boundStatement.setDecimal(0, (BigDecimal) value);
                break;
            case TIMESTAMP_WITHOUT_TIME_ZONE:
            case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
                boundStatement.setTimestamp(0, (Date) value);
                break;
            case BINARY:
            case VARBINARY:
                boundStatement.setBytes(0, (java.nio.ByteBuffer) value);
                break;
            case ARRAY:
                // 处理 Set 类型
                @SuppressWarnings("unchecked")
                Set<Object> setDataForDelete = (Set<Object>) value;
                boundStatement.setSet(0, setDataForDelete);
                break;
            default:
                boundStatement.setString(0, (String) value);
                break;
        }
        
        return boundStatement;
    }

    /**
     * 序列化 Set/Array 追加操作
     */
    public BoundStatement serializeAppendSet(RowData row, PreparedStatement preparedStatement, Set<String> appendSetFields) {
        BoundStatement boundStatement = preparedStatement.bind();
        
        List<String> fieldNames = rowType.getFieldNames();
        int paramIndex = 0;
        
        // 绑定需要追加的字段
        for (String fieldName : appendSetFields) {
            int fieldIndex = fieldNames.indexOf(fieldName);
            if (fieldIndex >= 0) {
                LogicalType fieldType = rowType.getTypeAt(fieldIndex);
                Object value = convertValue(row, fieldIndex, fieldType);
                if (fieldType.getTypeRoot() == ARRAY) {
                    // 对于 ARRAY 类型，转换为 Set 进行追加操作
                    boundStatement.setSet(paramIndex++, (Set<Object>) value);
                } else {
                    // 对于其他类型，直接设置值
                    setParameterValue(boundStatement, paramIndex++, value, fieldType);
                }
            }
        }
        
        // 绑定主键字段（用于 WHERE 子句）
        // 优先使用显式声明的主键，如果没有则使用除了 appendSetFields 之外的所有字段作为主键
        if (!primaryKeys.isEmpty()) {
            // 使用显式声明的主键
            for (String keyField : primaryKeys) {
                int fieldIndex = fieldNames.indexOf(keyField);
                if (fieldIndex >= 0) {
                    LogicalType keyFieldType = rowType.getTypeAt(fieldIndex);
                    Object keyValue = convertValue(row, fieldIndex, keyFieldType);
                    setParameterValue(boundStatement, paramIndex++, keyValue, keyFieldType);
                }
            }
        } else {
            // 回退到使用除了 appendSetFields 之外的所有字段作为主键
            for (int i = 0; i < fieldNames.size(); i++) {
                String fieldName = fieldNames.get(i);
                if (!appendSetFields.contains(fieldName)) {
                    LogicalType keyFieldType = rowType.getTypeAt(i);
                    Object keyValue = convertValue(row, i, keyFieldType);
                    setParameterValue(boundStatement, paramIndex++, keyValue, keyFieldType);
                }
            }
        }
        
        return boundStatement;
    }
    
    /**
     * 根据字段类型设置参数值
     */
    private void setParameterValue(BoundStatement boundStatement, int index, Object value, LogicalType fieldType) {
        if (value == null) {
            boundStatement.setToNull(index);
            return;
        }
        
        switch (fieldType.getTypeRoot()) {
            case CHAR:
            case VARCHAR:
                boundStatement.setString(index, (String) value);
                break;
            case INTEGER:
                boundStatement.setInt(index, (Integer) value);
                break;
            case BIGINT:
                boundStatement.setLong(index, (Long) value);
                break;
            case BOOLEAN:
                boundStatement.setBool(index, (Boolean) value);
                break;
            case FLOAT:
                boundStatement.setFloat(index, (Float) value);
                break;
            case DOUBLE:
                boundStatement.setDouble(index, (Double) value);
                break;
            case ARRAY:
                @SuppressWarnings("unchecked")
                Set<Object> setValue = (Set<Object>) value;
                boundStatement.setSet(index, setValue);
                break;
            default:
                boundStatement.setString(index, value.toString());
                break;
        }
    }

    private Object convertValue(RowData row, int fieldIndex, LogicalType fieldType) {
        if (row.isNullAt(fieldIndex)) {
            return null;
        }

        switch (fieldType.getTypeRoot()) {
            case CHAR:
            case VARCHAR:
                StringData stringData = row.getString(fieldIndex);
                return stringData != null ? stringData.toString() : null;
            case BOOLEAN:
                return row.getBoolean(fieldIndex);
            case TINYINT:
                return (int) row.getByte(fieldIndex);
            case SMALLINT:
                return (int) row.getShort(fieldIndex);
            case INTEGER:
                return row.getInt(fieldIndex);
            case BIGINT:
                return row.getLong(fieldIndex);
            case FLOAT:
                return row.getFloat(fieldIndex);
            case DOUBLE:
                return row.getDouble(fieldIndex);
            case DECIMAL:
                return row.getDecimal(fieldIndex, 0, 0);
            case TIMESTAMP_WITHOUT_TIME_ZONE:
                return Date.from(row.getTimestamp(fieldIndex, 3).toInstant());
            case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
                return Date.from(row.getTimestamp(fieldIndex, 3).toInstant());
            case BINARY:
            case VARBINARY:
                return java.nio.ByteBuffer.wrap(row.getBinary(fieldIndex));
            case ARRAY:
                // 处理 Set 类型（将 ArrayData 转换为 Set）
                ArrayData arrayData = row.getArray(fieldIndex);
                if (arrayData != null) {
                    ArrayType arrayType = (ArrayType) fieldType;
                    LogicalType elementType = arrayType.getElementType();
                    
                    Set<Object> setData = new HashSet<>();
                    for (int i = 0; i < arrayData.size(); i++) {
                        if (elementType.getTypeRoot() == VARCHAR || 
                            elementType.getTypeRoot() == CHAR) {
                            // 处理 set<text> 类型
                            StringData stringElement = arrayData.getString(i);
                            if (stringElement != null) {
                                setData.add(stringElement.toString());
                            }
                        } else {
                            // 处理其他元素类型，这里简化处理，主要支持 set<text>
                            // 对于其他类型，可以根据需要扩展
                            StringData stringElement = arrayData.getString(i);
                            if (stringElement != null) {
                                setData.add(stringElement.toString());
                            }
                        }
                    }
                    return setData;
                }
                return null;
            default:
                // 对于其他类型，转换为字符串
                return row.getString(fieldIndex).toString();
        }
    }
} 